airbyte/_util/api_util.py (1)

2095-2135: Clean addition following existing patterns.

The function mirrors the structure of get_connection_state nicely. One thought: the docstring mentions HTTP 423 behavior, but _make_config_api_request will surface it as a generic AirbyteError. Would it be helpful to catch 423 specifically and raise a more descriptive error (e.g., "Cannot update state while a sync is running") so users get actionable feedback, wdyt?

airbyte/cloud/connections.py (2)

453-478: set_state doesn't validate that connection_state contains a matching connectionId — could this lead to silent mismatches?

The docstring says connectionId inside the dict "Must match this connection's ID," but there's no enforcement. If a caller passes a state dict from a different connection (e.g., copy-paste error), the API would receive a top-level connectionId (from self.connection_id) that might conflict with the one inside connection_state. Depending on which one the API respects, this could either silently do the right thing or corrupt another connection's state. A quick guard here could prevent a painful debugging session, wdyt?

🛡️ Proposed validation

     def set_state(
         self,
         connection_state: dict[str, Any],
     ) -> ConnectionStateResponse:
+        if (
+            "connectionId" in connection_state
+            and connection_state["connectionId"] != self.connection_id
+        ):
+            raise PyAirbyteInputError(
+                message="connectionId in state payload does not match this connection.",
+                context={
+                    "expected": self.connection_id,
+                    "actual": connection_state["connectionId"],
+                },
+            )
         updated_data = api_util.create_or_update_connection_state_safe(

---

504-583: set_stream_state uses a fetch-modify-push pattern — inherent TOCTOU, but acceptable given the safe endpoint.

The read-modify-write on lines 529→583 means two concurrent set_stream_state calls could race, with the last writer winning. The "safe" endpoint guards against sync-time corruption (423), but not against concurrent client updates. This is likely acceptable for the use case, but maybe worth a brief note in the docstring so callers are aware, wdyt?

airbyte/_util/api_util.py (1)

2095-2142: Clean implementation, consistent with existing patterns.

The docstring does a great job explaining the full-replacement semantics and the HTTP 423 behavior. One minor thought: the connection_state dict likely contains its own connectionId field, and there's no validation that it matches the connection_id parameter passed to this function. A mismatch could lead to confusing API errors. Would it be worth adding a quick sanity check, or is that better left to the caller, wdyt?

💡 Optional validation

+    # Sanity check: if connectionState includes a connectionId, it should match
+    inner_conn_id = connection_state.get("connectionId")
+    if inner_conn_id and inner_conn_id != connection_id:
+        raise PyAirbyteInputError(
+            message="connectionId in connection_state does not match the connection_id parameter.",
+            context={
+                "connection_id": connection_id,
+                "connection_state_connection_id": inner_conn_id,
+            },
+        )
     return _make_config_api_request(
         path="/state/create_or_update_safe",

airbyte/_util/api_util.py (1)

38-43: Duplicate TYPE_CHECKING block — could you merge these two?

Lines 38-40 and 42-47 both guard if TYPE_CHECKING: imports. The first one imports Callable and the second imports Callable again (plus DestinationConfiguration). This results in a redundant import of Callable and two separate blocks that could be consolidated, wdyt?

🔧 Proposed fix

 if TYPE_CHECKING:
     from collections.abc import Callable
-
-
-if TYPE_CHECKING:
-    from collections.abc import Callable
 
     from airbyte_api.models import (
         DestinationConfiguration,
     )

airbyte/_util/api_util.py (1)

1467-1472: No timeout on the requests.request() call — could block indefinitely.

_make_config_api_request doesn't pass a timeout to requests.request(). If the config API hangs, callers (including the new state/catalog methods) would block forever. Would adding a reasonable default timeout (e.g., timeout=60) be worth considering here, wdyt?

🔧 Proposed fix

     response = requests.request(
         method="POST",
         url=full_url,
         headers=headers,
         json=json,
+        timeout=60,
     )